tree = {}
tree_mt = {__index = tree}

--*************** construction **************
function tree:new( info , cover, meaningful )
	local t = {}

	-- extracting infor from label (e.g. S[ate\VBD]=H )
	if type(info) == 'table' then
		t.label = info.label
		t.hWord = info.hWord
		t.isHead = info.isHead
	else
		local i,j = info:find('.+%[') 
		if i ~= nil and j ~= nil then 
			t.label = info:sub(i,j-1 )
			i,j = info:find('%[.+%/'); t.hWord = info:sub(i+1,j-1)
			i,j = info:find('=H'); t.isHead = true; if i == nil then t.isHead = false end
		else t.label = info end
	end

	t.children = {}
	t.cover = cover or {}

	t.match = nil
	t.meaningful = meaningful or true

	setmetatable(t, tree_mt)
	return t
end

function tree:createFromString( input , leafId )
	local str = input:sub(2,input:len()-1)
	local leafId = leafId or 1

	-- read label
	local i = str:find(' ') --print(input)
	local t = tree:new( str:sub(1,i-1) )

	-- read children
	i = i + 1
	
	-- check if this is POS
	if str:sub(i,i) ~= '(' then
		t.children[1] = tree:new( str:sub(i)  )
		t.children[1].cover = {leafId, leafId}
		t.cover = {leafId, leafId}
		leafId = leafId + 1

	else
		local countOpen = 1
		str = str:sub(i)
		local j = 2
		while true do
			local c = str:sub(j,j)
			if c == '(' then 
				countOpen = countOpen + 1
			elseif c == ')' then
				countOpen = countOpen - 1
			end

			if countOpen == 0 then
				t.children[#t.children+1],leafId = tree:createFromString( str:sub(1,j) , leafId) 
				if j + 1 > str:len() then break end
				j = str:find('(',j+1,true)
				str = str:sub(j)
				countOpen = 1
				j = 2
			else 
				j = j + 1
			end
		end
		t.cover = {t.children[1].cover[1], t.children[#t.children].cover[2]}
	end

	return t, leafId
end

--******************** deep copy ****************
function tree:clone()
	local nChild = #self.children
	local newtree = tree:new( {self.label, self.isHead, self.hWord} , self.cover , self.meaningful )
	
	-- clone children
	for i = 1,nChild do
		newtree.children[i] = self.children[i]:clone()
	end

	return newtree
end

--***************** toString ****************
function tree:toString( onlyLeaves ) 
	local onlyLeaves = onlyLeaves or false print(self)

	local str = self.label .. '[' .. (self.hWord or '') .. ']'
	if self.isHead then str = str .. '=H' end
	if self.meaningful == false then str = str .. '*' end

	local nChild = #self.children
	if nChild == 0 then
		if onlyLeaves then 
			return str .. ' '
		else 
			return str
		end
	else
		if onlyLeaves == false then 
			str = '(' .. str 
		else 
			str = ''
		end
		for i = 1,nChild do
			if onlyLeaves == false then 
				str = str .. ' ' .. self.children[i]:toString(onlyLeaves)
			else 
				str = str .. self.children[i]:toString(onlyLeaves) 
			end
		end
		if onlyLeaves == false then str = str .. ')' end
	end
	return str
end

--******************* extract leaves ***************
function tree:leaves( ret )
	local ret = ret or {}
	local nChild = #self.children
	if nChild == 0 then
		ret[#ret+1] = self
	else
		for i = 1,nChild do
			self.children[i]:leaves(ret)
		end
	end
	return ret
end

--********************** get all nodes ******************
-- pre-order traverse
function tree:allNodes( ret )
	local ret = ret or {}
	local nChild = #self.children
	ret[#ret+1] = self

	for i = 1,nChild do
		self.children[i]:allNodes( ret )
	end
	return ret
end

--******************* extract compact tree *****************
-- pre-order traverse
function tree:compact( ret )
	local ret = ret or {}
	local id = #ret + 1
	local t = {
		hWord = self.hWord,
		meaningful = self.meaningful, 
		childId = {} }

	if #self.children == 0 then t.label = self.label end	

	ret[id] = t
	for i = 1,#self.children do
		ret, cid = self.children[i]:compact( ret )
		t.childId[i] = cid
	end
	return ret, id
end

function tree.comp2string( comp , str, id )
	local str = str or ''
	local id = id or 1
	local node = comp[id]
	str = str .. ' '
	if #node.childId == 0 then str = str .. node.label
	else
		str = str .. '(' .. (node.label or 'X') .. '[' .. node.hWord .. ']'
		if node.isHead then str = str .. '=H' end
		for _,cid in ipairs(node.childId) do
			str = tree.comp2string( comp, str, cid )
		end
		str = str .. ')'
	end
	return str
end

--****************** load forest from file ****************
function tree:loadForestFromText( filename )
	local forest = {}
	for line in io.lines(filename) do
		forest[#forest+1] = tree:createFromString(line)
	end
	return forest
end

function tree.loadForestFromBin( filename )
	local f = torch.DiskFile(filename, 'r')
	f:binary()
	local forest = f:readObject()
	f:close()

	print('set metatable')
	local function makeTree ( t )
		setmetatable(t, tree_mt)
		for i = 1,#t.children do
			makeTree( t.children[i] )
		end
	end

	for i=1,#forest do
		makeTree(forest[i])
	end

	return forest
end

--********************** binarize ********************
-- if strict == true then nChild == 2
-- else nChild <= 2
function tree:binarize( strict ) 
	local strict = strict or false
	local nChild = #self.children
	
	if strict then
		while nChild == 1 do
			self.label = self.label .. '^' .. self.children[1].label 
			self.children = self.children[1].children
			nChild = #self.children
		end
	end

	if nChild > 2 then
		-- newlabel = '@[this_node_label]'
		local newlabel = self.label
		if newlabel:sub(1,1) ~= '@' then newlabel = '@' .. newlabel end

		-- create new child
		local lChild, rChild
		if self.children[1].isHead then 
			rChild = self.children[nChild]
			lChild = tree:new( { label = newlabel, isHead = true, hWord = self.hWord } )
			for i = 1,nChild-1 do
				lChild.children[i] = self.children[i]
			end
			lChild.cover[1] = lChild.children[1].cover[1]
			lChild.cover[2] = lChild.children[nChild-1].cover[2]
		else
			lChild = self.children[1]
			rChild = tree:new( { label = newlabel, isHead = true, hWord = self.hWord } )
			for i = 2, nChild do 
				rChild.children[i-1] = self.children[i]
			end
			rChild.cover[1] = rChild.children[1].cover[1]
			rChild.cover[2] = rChild.children[nChild-1].cover[2]
		end
	
		-- update child list for this node
		self.children = {lChild, rChild}
	end
	
	-- binarize subtrees
	for i = 1,#self.children do
		self.children[i]:binarize(strict)
	end
end

--*********************** remove labels *************************--
function tree:removeTags()
	local nChild = #self.children

	if nChild == 0 then
		-- keep words at leaves
		local newlabel
		for tok in self.label:gmatch('[^^]+') do
			newlabel = tok
		end
		self.label = newlabel
	else
		self.label = 'X'
		for i = 1,nChild do
			self.children[i]:removeTags()
		end
	end
end

--****************** compare from leaves *****************
function tree:compare( goldTree )
	local Nodes = self:allNodes()
	local GNodes = goldTree:allNodes()
	
	local nNodes = #Nodes
	local nGNodes = #GNodes

	for i = nNodes,1,-1 do
		local node = Nodes[i]
		local nChild = #node.children
		node.meaningful = false

		for j = nGNodes,1,-1 do
			local gnode = GNodes[j]
			local nGChild = #gnode.children

			if node.cover[1] == gnode.cover[1] and node.cover[2] == gnode.cover[2] and 
						nChild == nGChild and node.label == gnode.label and
						node.isHead == gnode.isHead and node.hWord == gnode.hWord then 
				local ok = true
				for ii = 1,nChild do
					if node.children[ii].match ~= gnode.children[ii] then
						ok = false
						break
					end
				end
				if ok then 
					node.match = gnode 
					node.meaningful = true
					break 
				end
			end
		end
	end
end

--**************** for testing ***************
function testLoadForest()
	local filename = '/media/DATA/Wind/PhD/NeuralNetwork/semreranker/Data/WSJ_wo_trace/train'
	local forest = tree:loadForest(filename)
	
	-- print sentences
	for _,tree in ipairs(forest) do
		local leaves = tree:leaves()
		local str = ''
		for _,node in ipairs(leaves) do
			str = str .. node.label .. ' '
		end
		print(str)
	end
end

--testLoadForest()

function test()
	local t = tree:createFromString('(ROOT[ate/VBD] (S[ate/VBD]=H (NP[teacher/NN] (DT[The/DT] The) (NN[teacher/NN]=H teacher)) (VP[ate/VBD]=H (VBD[ate/VBD]=H ate) (NP[spaghetti/NNS] (DT[the/DT] the) (NNS[spaghetti/NNS]=H spaghetti))) (.[./.] .)))')
	t:binarize(true)
	print(t:toString())
	t:removeTags()
	local ct = t:compact()
	print(tree.comp2string(ct))
end
--test()
